import com.rwind.model.Action;
import com.rwind.model.ActionPattern;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
import org.apache.flink.util.Collector;

public class BroadcastStateExample {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        // 读取用户行为事件流
        DataStreamSource<Action> actionStream =
                env.fromElements(new Action("Alice", "login"),
                        new Action("Alice", "pay"), new Action("Bob", "login"), new Action("Bob", "buy")
                );


        // 定义行为模式流，代表了要检测的标准
        DataStreamSource<ActionPattern> patternStream = env
                .fromElements(
                        new ActionPattern("login", "pay"), new ActionPattern("login", "buy")
                );

        // 定义广播状态的描述器，创建广播流
        MapStateDescriptor<Void, ActionPattern> bcStateDescriptor = new MapStateDescriptor<>("patterns", Types.VOID, Types.POJO(ActionPattern.class));
        BroadcastStream<ActionPattern> bcPatterns = patternStream.broadcast(bcStateDescriptor);

        // 将事件流和广播流连接起来，进行处理
        DataStream<Tuple2<String, ActionPattern>> matches = actionStream
                .keyBy(data -> data.userId)
                .connect(bcPatterns)
                .process(new PatternEvaluator());
        matches.print();

        env.execute();


    }

    public static class PatternEvaluator extends KeyedBroadcastProcessFunction<String, Action, ActionPattern, Tuple2<String, ActionPattern>> {

        // 定义一个值状态，保存上一次用户行为
        ValueState<String> prevActionState;

        @Override
        public void open(Configuration conf) {
            prevActionState = getRuntimeContext().getState(
                    new ValueStateDescriptor<>("lastAction", Types.STRING));
        }


        @Override
        public void processElement(Action action, KeyedBroadcastProcessFunction<String, Action, ActionPattern, Tuple2<String, ActionPattern>>.ReadOnlyContext ctx,
                                   Collector<Tuple2<String, ActionPattern>> collector) throws Exception {
            ActionPattern pattern = ctx.getBroadcastState(new MapStateDescriptor<>("patterns", Types.VOID, Types.POJO(ActionPattern.class))).get(null);
            String prevAction = prevActionState.value();
            if (pattern != null && prevAction != null) {
// 如果前后两次行为都符合模式定义，输出一组匹配
                if (pattern.action1.equals(prevAction)
                        && pattern.action2.equals(action.action)) {
                    collector.collect(new Tuple2<>(ctx.getCurrentKey(), pattern));
                }
            }
// 更新状态
            prevActionState.update(action.action);
        }

        @Override
        public void processBroadcastElement(ActionPattern actionPattern,
                                            KeyedBroadcastProcessFunction<String, Action, ActionPattern, Tuple2<String, ActionPattern>>.Context context,
                                            Collector<Tuple2<String, ActionPattern>> collector) throws Exception {
            BroadcastState<Void, ActionPattern> bcState = context.getBroadcastState(
                    new MapStateDescriptor<>("patterns", Types.VOID, Types.POJO(ActionPattern.class)));
            // 将广播状态更新为当前的 pattern
            bcState.put(null, actionPattern);
        }
    }
}
